

from flask import Flask, Response, stream_with_context
import cv2
import tensorflow as tf
from tf_pose.estimator import TfPoseEstimator
from tf_pose.networks import get_graph_path, model_wh
from tf_pose import common
from pose_estimate import PoseEstimate
import numpy as np
app = Flask(__name__)


def process_frame(frame, e, pe, timeF, data_sequence, result, w, h):
    # 估计姿态
    humans = e.inference(frame, resize_to_default=(w > 0 and h > 0), upsample_size=4.0)

    # 在帧上绘制人体骨骼
    frame = TfPoseEstimator.draw_humans(frame, humans, imgcopy=False)

    # 提取人体关键点信息
    centers = {}
    flat = [0.0 for i in range(36)]
    image_h, image_w = frame.shape[:2]
    for human in humans:
        for i in range(common.CocoPart.Background.value):
            if i not in human.body_parts.keys():
                continue
            body_part = human.body_parts[i]
            center = (int(body_part.x * image_w), int(body_part.y * image_h))
            centers[i] = center
            flat[i * 2] = center[0]
            flat[i * 2 + 1] = center[1]
            cv2.circle(frame, center, 3, common.CocoColors[i], thickness=3, lineType=8, shift=0)
            text = "X:{} Y:{}".format(center[0], center[1])
            cv2.putText(frame, text, (center[0], center[1]), cv2.FONT_HERSHEY_PLAIN, 1, (255, 255, 255), thickness=1)

            # 序列化数据并添加到数据序列中
    if len(humans) > 0:
        fame_feature = []
        human_body = humans[0].body_parts
        for i in range(17):
            body_feature = []
            try:
                body_feature.append(human_body[i].part_idx)
                body_feature.append(human_body[i].x)
                body_feature.append(human_body[i].y)
                body_feature.append(human_body[i].score)
            except:
                body_feature.append(i)
                body_feature.append(0)
                body_feature.append(0)
                body_feature.append(0)
            fame_feature.append(body_feature)
        arr = np.array(fame_feature)
        arr = np.reshape(arr, [-1])
        data_sequence.append(arr)

        # 如果数据序列长度达到阈值，进行姿态分类并重置数据序列
    if len(data_sequence) == timeF:
        input_data = np.array(data_sequence)
        input_data = input_data.reshape([1, timeF, 68])
        pre = pe.estimate(input_data)
        print(pre)
        result = pre[0][0][1]
        data_sequence = []

        # 在帧上绘制分类结果
    cv2.putText(frame, "{:.4f} {}".format(result, "standard"), (0, 50), cv2.FONT_HERSHEY_PLAIN, 3.5, (0, 0, 255), 3)

    return frame, result, data_sequence


@app.route('/tfposecamera')
def supvisedcamera():
    def generate():
        # 初始化视频捕获和TfPoseEstimator
        videoCapture = cv2.VideoCapture(0)  # 从摄像头捕获视频
        fps = videoCapture.get(cv2.CAP_PROP_FPS)
        size = (int(videoCapture.get(cv2.CAP_PROP_FRAME_WIDTH)), int(videoCapture.get(cv2.CAP_PROP_FRAME_HEIGHT)))

        # 初始化姿态估计器
        resolution = "432x368"
        model = "cmu"
        w, h = model_wh(resolution)
        config = tf.ConfigProto(log_device_placement=True)
        with tf.Graph().as_default():
            e = TfPoseEstimator(get_graph_path(model), target_size=(w, h), tf_config=config)

            # 初始化姿态分类器
        pe = PoseEstimate("./feed_lstm_attention-fwcce-1/model")

        # 设置序列化和分类相关参数
        timeF = 10
        data_sequence = []
        result = 0

        while True:
            ret, frame = videoCapture.read()
            if not ret:
                break

                # 处理帧
            frame, result, data_sequence = process_frame(frame, e, pe, timeF, data_sequence, result, w, h)

            # 将图像编码为 JPEG 格式并通过网络传输
            ret, buffer = cv2.imencode('.jpg', frame)
            if not ret:
                break

                # 检查退出键
            key = cv2.waitKey(1)  # 用于处理视频或实时摄像头捕获的帧
            if key == 27:  # 按 ESC 键退出
                break

            yield (b'--frame\r\n'
                   b'Content-Type: image/jpeg\r\n\r\n' + buffer.tobytes() + b'\r\n')

    return Response(stream_with_context(generate()), mimetype='multipart/x-mixed-replace; boundary=frame')


if __name__ == "__main__":
    app.run(host="0.0.0.0", port=10001, debug=True)